package top.atcra.rpc.protection.rateLimit.impl;

import top.atcra.rpc.protection.rateLimit.RateLimiter;

public class TokenBucketLimiter implements RateLimiter {
    private final int capacity;
    private final int tokensPerSecond;
    private int tokenCount;
    private long lastRefillTime;

    public TokenBucketLimiter(int capacity, int tokensPerSecond) {
        this.capacity = capacity;
        this.tokensPerSecond = tokensPerSecond;
        this.lastRefillTime = System.currentTimeMillis();
        tokenCount = capacity;
    }

    @Override
    public synchronized boolean allowEnter() {
        long currentTime = System.currentTimeMillis();
        long elapsedTime = currentTime - lastRefillTime;

        int tokensToAdd = (int) (elapsedTime * tokensPerSecond / 1000);
        tokenCount = Math.min(tokenCount + tokensToAdd, capacity);
        lastRefillTime = currentTime;

        if (tokenCount > 0) {
            tokenCount--;
            return true;
        }

        return false;
    }
}
